#!/usr/bin/python
# -*-coding:utf-8-*-
import os
import pandas as pd
import numpy as np

### 底层读取数据的依赖（不提供）
from zbc_factor_lib.base.factors_library_base import NewRQFactorLib as DataReader

from dir_info import ylabel_data_dir
from dir_info import cache_data_dir

from base.utils import get_next_trade_date

data_reader = DataReader()

# data_reader.show_basic_library_db()

def get_n_day_return_label_main(n_days=60):
    ##
    processed_trade_date_data = data_reader.read_basic_data_table('processed_trade_date_data')

    processed_trade_date_data = processed_trade_date_data.set_index('natural_date')

    processed_trade_date_data['start'] = processed_trade_date_data['trade_date']
    processed_trade_date_data['end'] = processed_trade_date_data['trade_date'].shift(-n_days)

    # TODO - 读取label日期
    target_construction_label = pd.read_hdf(os.path.join(cache_data_dir, 'X_financial_statement_data_target_construction_label.h5'))

    target_construction_label['target_start'] = processed_trade_date_data.loc[target_construction_label['declare_date'], 'start'].values
    target_construction_label['target_end'] = processed_trade_date_data.loc[target_construction_label['declare_date'], 'end'].values

    target_construction_label = target_construction_label.reset_index()

    # TODO - 读取数据
    processed_daily_stock_trade_data = data_reader.read_basic_data_table('processed_daily_stock_trade_data')

    processed_daily_stock_trade_data = processed_daily_stock_trade_data.rename(columns={'code': 'stock_code'})

    processed_daily_stock_trade_data['bclose'] = \
        processed_daily_stock_trade_data['close'] * processed_daily_stock_trade_data['back_adjfactor']

    processed_daily_stock_trade_data = processed_daily_stock_trade_data.set_index(['date', 'stock_code'])

    # map
    map_keys = target_construction_label[['target_start', 'stock_code']].set_index(['target_start', 'stock_code']).index

    target_construction_label['start_price'] = processed_daily_stock_trade_data.loc[map_keys, 'bclose'].values

    map_keys = target_construction_label[['target_end', 'stock_code']].set_index(['target_end', 'stock_code']).index

    target_construction_label['end_price'] = processed_daily_stock_trade_data.loc[map_keys, 'bclose'].values

    target_construction_label['y'] = target_construction_label['end_price'] / target_construction_label['start_price'] - 1

    target_construction_label = target_construction_label.set_index(['report_date', 'stock_code'])

    target_construction_label = target_construction_label[['y']]

    target_construction_label = target_construction_label.astype(np.float32)

    # TODO - save
    save_filename = 'financial_statement_y_%dd_return_label' % (n_days)

    target_construction_label.to_hdf(os.path.join(ylabel_data_dir, save_filename+'.h5'),
                                     key=save_filename,
                                     mode='w',
                                     format='table')

    print(save_filename, 'done!\n')


if __name__ == '__main__':
    n_days_list = [30, 60, 90]

    for n_days in n_days_list:
        get_n_day_return_label_main(n_days=n_days)

